package org.ragdollcat.secondaiagent.agent;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.ObjUtil;
import cn.hutool.core.util.StrUtil;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.ragdollcat.secondaiagent.agent.enums.AgentState;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;

/**
 * agent代理基础层
 */
@Data
@Slf4j
public abstract class BaseAgent {

    /**
     * 名称
     */
    private String name;

    /**
     * 系统提示词
     */
    private String systemPrompt;

    /**
     * 下一步的提示词
     */
    private String nextStepPrompt;


    /**
     * 状态，默认是空闲状态
     */
    private AgentState agentState = AgentState.IDLE;

    /**
     * 当前步骤，用于多轮推论
     */
    private Integer currentStep = 1;

    /**
     * 最大轮次
     */
    private Integer maxStep = 10;


    /**
     * 大模型客户端
     */
    private ChatClient chatClient;


    /**
     * 自主维护上下文记忆
     */
    private List<Message> memoryMessages = new ArrayList<>();

    /**
     * 智能体循环检测阈值
     */
    private Integer duplicateThreshold = 2;

    /**
     * 运行代理
     *
     * @param text
     * @return
     */
    public String run(String text) {
        if (!agentState.equals(AgentState.IDLE)) {
            throw new RuntimeException("当前代理 非空闲状态！");
        }
        if (StrUtil.isBlank(text)) {
            throw new RuntimeException("用户输入信息不能为空！");
        }
        //修改状态
        agentState = AgentState.RUNNING;
        //记录上下文信息(用户提示词)
        memoryMessages.add(new UserMessage(text));

        ArrayList<String> list = new ArrayList<>();
        try {
            //循环条件 小于最大轮次，并且状态不为结束
            for (int i = 0; i < maxStep && agentState != AgentState.FINISHED; i++) {
                //执行具体的操作,子类实现
                String stepResult = step(null);


                if (this.isStuck()){
                    this.handleStuckState();
                }


                stepResult = stepResult + "当前轮次" + currentStep;
                log.info("当前轮次：{}", currentStep);
                //当前轮次 + 1
                currentStep = currentStep + 1;
                list.add(stepResult);
                //检查是否超过了步数限制
                if (currentStep >= maxStep) {
                    agentState = AgentState.FINISHED;
                    list.add("终止，到达最大步数(" + maxStep + ")");
                }
            }
            return String.join("/n", list);
        } catch (Exception e) {
            agentState = AgentState.ERROR;
            log.error("执行错误，原因:", e);
            return "执行错误" + e.getMessage();
        } finally {
            this.clearUp();
        }


    }


    /**
     * 运行代理
     *
     * @param text
     * @return
     */
    public SseEmitter runAsync(String text) {

        SseEmitter sseEmitter = new SseEmitter(300000L);

        //先让SseEmitter返回，然后异步执行
        CompletableFuture.runAsync(() -> {
            try {
                if (!agentState.equals(AgentState.IDLE)) {
                    sseEmitter.send("当前代理 非空闲状态！");
                    sseEmitter.complete();
                }
                if (StrUtil.isBlank(text)) {
                    sseEmitter.send("用户输入信息不能为空！");
                    sseEmitter.complete();
                }
                //修改状态
                agentState = AgentState.RUNNING;
                //记录上下文信息(用户提示词)
                memoryMessages.add(new UserMessage(text));

                try {
                    //循环条件 小于最大轮次，并且状态不为结束
                    for (int i = 0; i < maxStep && agentState != AgentState.FINISHED; i++) {
                        //执行具体的操作,子类实现
                        String stepResult = step(sseEmitter);

                        if (this.isStuck()){
                            this.handleStuckState();
                        }

                        log.info("当前轮次：{}", currentStep);
                        //当前轮次 + 1
                        currentStep = currentStep + 1;
                        //每一步完成后都发送消息
                        sseEmitter.send(stepResult);

                        //检查是否超过了步数限制
                        if (currentStep >= maxStep) {
                            agentState = AgentState.FINISHED;
                            //超出最大步数限制，也发送消息
                            sseEmitter.send("终止，到达最大步数(" + maxStep + ")");
                        }
                    }
                    //最终完成
                    sseEmitter.complete();
                } catch (Exception e) {
                    agentState = AgentState.ERROR;
                    log.error("执行错误，原因:", e);
                    try {
                        sseEmitter.send("执行错误" + e.getMessage());
                        sseEmitter.complete();
                    } catch (IOException ex) {
                        sseEmitter.completeWithError(ex);
                    }
                } finally {
                    this.clearUp();
                }
            } catch (Exception e) {
                sseEmitter.completeWithError(e);
            }
        });

        //设置超时和完成回调
        sseEmitter.onTimeout(()->{
            this.agentState = AgentState.ERROR;
            this.clearUp();
            log.error("SSE execute time out");
        });

        sseEmitter.onCompletion(()->{
            if (this.agentState == AgentState.RUNNING) {
                this.agentState = AgentState.FINISHED;
            }
            this.clearUp();
            log.info("SSE completed");
        });


        return sseEmitter;

    }

    /**
     * 循环消息检测
     * @return
     */
    public boolean isStuck() {
//        检查消息数量是否足够进行循环检测（至少需要2条消息）
        if (this.memoryMessages.size() < 2) {
            return false;
        }
        // 获取最后一条消息（最新的助手回复）
        Message lastMessage = CollUtil.getLast(this.memoryMessages);
        if (ObjUtil.isEmpty(lastMessage) || lastMessage.getText().isEmpty()){
            return false;
        }

        //统计与最后一条消息内容相同的助手消息数量
        int count = 0;
        //遍历除最后一条外的所有历史消息（倒序）
        for (int i = this.memoryMessages.size() - 2 ; i >=  0; i--) {
            Message message = memoryMessages.get(i);
//            只统计助手角色且内容完全相同的消息
            if (message instanceof AssistantMessage && message.getText().equals(lastMessage.getText())){
                count++;
            }
        }
//        判断重复次数是否达到阈值，如果达到则认为智能体陷入循环
        return count >= duplicateThreshold;
    }


    /**
     * 处理循环消息
     */
    public void handleStuckState(){
        String stuckPrompt = "Observed duplicate responses. Consider new strategies and avoid repeating ineffective paths already attempted.";
        this.nextStepPrompt = nextStepPrompt + "/n" + stuckPrompt;
        log.info("Agent detected stuck state. Added prompt: {}",stuckPrompt);
    }


    /**
     * 定义每次循环中执行的步骤，子类实现
     *
     * @return
     */
    public abstract String step(SseEmitter sseEmitter);


    /**
     * 清理资源
     */
    public void clearUp() {
        log.info("清理资源....");
    }

}

